from llmweb.base import LLMWebBase, LoginMixin
from llmweb import BrowserPool
from typing import Optional,AsyncGenerator
from playwright.async_api import Locator, Page
from llmweb.logging import logger
import time
import asyncio
import re

async def keep_browser_open(page, timeout):
    # 这个函数会等待指定的时间，直到超时或被取消
    await asyncio.sleep(timeout)
class KimLLM(LLMWebBase, LoginMixin):
    wait_timeout:int = 3000 # seconds
    def __init__(self, pool: Optional[BrowserPool]=BrowserPool()):
        site = "https://kimi.moonshot.cn/"
        super().__init__(site, pool) 
    
    async def generate(self, prompt: str, **kwargs) -> AsyncGenerator[str, None]:
        page: Page = await self.pool.acquire_page()
        await page.route(re.compile(r"\.(jpg|png|svg)$"), lambda route: route.abort()) 
        await page.goto(self.site)
        stop = kwargs.get("stop")
        logger.debug(f"input Stop: {stop}")
        try:
            await self.pool.load_local_storage(page)
            await self.pool.load_cookies(page)
            await page.reload()
        except Exception as e:
            logger.error(f"Load local storage error: {e}")
            await self.pool.release_page(page)
            raise e
        
        await page.wait_for_load_state('networkidle')
        await page.get_by_test_id("msh-chatinput-editor").fill(prompt)
        await page.get_by_test_id("msh-chatinput-send-button").click()
        
        await page.get_by_text("停止输出").wait_for(timeout=self.wait_timeout)
        text_locator: Locator = page.locator('div[class^="markdown___"]')

        total_timeout_time = time.time() + self.wait_timeout
        chunk_wait_timeout = time.time() + 5 # 5秒
        text = ""
        while True:
            try:
                if time.time() > total_timeout_time:
                        ## 总体超时，退出循环
                        raise TimeoutError("text response timeout")
                
                els = await text_locator.all()
                n = len(els)
                if n == 1:
                    continue
                last = els[n-1]
                new_text = await last.inner_text()
                if stop is not None and isinstance(stop,list):
                        for s in stop:
                            if s in text:
                                await page.get_by_text("停止输出").click()
                                logger.debug("stop text")
                                raise StopIteration(s)
                if text != new_text:
                    chunks = new_text.replace(text, "")
                    yield chunks
                    text = new_text
                    chunk_wait_timeout = time.time() + 5
                else:
                    if time.time() > chunk_wait_timeout:
                        #等待单个块超时，退出循环
                        break
            except Exception as e:
                logger.error(e)
                break        
            await asyncio.sleep(0.5) 

        await self.pool.release_page(page)
        logger.debug("exit kimi generator")
    async def login(self, **kwargs):
        """
        需要人工登录，验证，存储登录信息
        """
        page = await self.pool.acquire_page()
        await page.goto(self.site)
        await page.get_by_text("登录").click()
        page.on('load', lambda: asyncio.create_task(self.pool.add_save_cookie_btn(page)))
        await asyncio.sleep(30)
        await self.pool.release_page(page)
